#include "stdafx.h"
#include "modelLoader.h"
#include "meshGeo.h"
#include "DDSTextureLoader.h"
#include "../frameresource.h"

ModelLoader::ModelLoader(Renderer* ptr, const XMMATRIX& mat): pRender(ptr), model_matrix(mat)
{

}
ModelLoader::~ModelLoader()
{

}

bool ModelLoader::loadModel(const std::string& pFile, UINT& objIdx, UINT& matCBIndex, UINT& texIndex)
{
	m_objIdx = objIdx;
	m_matCBIndex = matCBIndex;
	m_texIndex = texIndex;
	int i= pFile.size() - 1;
	for (; i >= 0; i--)
		if (pFile[i] == '/')
			break;
	directory = pFile.substr(0,i);
	// Create an instance of the Importer class
	Assimp::Importer importer;
	// And have it read the given file with some example postprocessing
	// Usually - if speed is not the most important aspect for you - you'll
	// probably to request more postprocessing than we do in this example.
	const aiScene* scene = importer.ReadFile(pFile,
		//aiProcess_CalcTangentSpace |
		aiProcess_Triangulate |
		aiProcess_JoinIdenticalVertices |
		//aiProcess_SortByPType |
		aiProcess_ConvertToLeftHanded);
	// If the import failed, report it
	if (!scene)
	{
		std::cout << importer.GetErrorString() << std::endl;
		return false;
	}

	geo = std::make_unique<MeshGeometry>();
	geo->Name = "model";
	model_name = geo->Name;

	processNode(scene->mRootNode, scene);

	vertices.shrink_to_fit();
	indices.shrink_to_fit();

	const UINT vbByteSize = (UINT)vertices.size() * sizeof(Vertex);
	const UINT ibByteSize = (UINT)indices.size() * sizeof(std::uint16_t);
	
	ThrowIfFailed(D3DCreateBlob(vbByteSize, &geo->VertexBufferCPU));
	CopyMemory(geo->VertexBufferCPU->GetBufferPointer(), vertices.data(), vbByteSize);

	ThrowIfFailed(D3DCreateBlob(ibByteSize, &geo->IndexBufferCPU));
	CopyMemory(geo->IndexBufferCPU->GetBufferPointer(), indices.data(), ibByteSize);

	geo->VertexBufferGPU = D3DUtil::CreateDefaultBuffer(pRender->m_device.Get(),
		pRender->m_commandList.Get(), vertices.data(), vbByteSize, geo->VertexBufferUploader,
		D3D12_RESOURCE_STATE_VERTEX_AND_CONSTANT_BUFFER);
	geo->IndexBufferGPU = D3DUtil::CreateDefaultBuffer(pRender->m_device.Get(),
		pRender->m_commandList.Get(), indices.data(), ibByteSize, geo->IndexBufferUploader,
		D3D12_RESOURCE_STATE_INDEX_BUFFER);

	geo->VertexByteStride = sizeof(Vertex);
	geo->VertexBufferByteSize = vbByteSize;
	geo->IndexFormat = DXGI_FORMAT_R16_UINT;
	geo->IndexBufferByteSize = ibByteSize;

	std::cout << "model vertices : " << vertices.size() << std::endl;
	std::cout << "model indices : " << indices.size() << std::endl;

	pRender->m_meshes[geo->Name] = std::move(geo);

	objIdx = m_objIdx;
	matCBIndex = m_matCBIndex;
	texIndex = m_texIndex;

	return true;
}

void ModelLoader::processMesh(aiMesh* mesh, const aiScene* scene)
{
	static UINT idx = 0;
	SubmeshGeometry submesh;
	submesh.BaseVertexLocation = vertices.size();
	submesh.StartIndexLocation = indices.size();
	
	// Data to fill
	if (mesh->mMaterialIndex >= 0)
	{
		aiMaterial* mat = scene->mMaterials[mesh->mMaterialIndex];
		if (textype.empty())
		{
			textype = determineTextureType(scene, mat);
		}	
	}


	// Walk through each of the mesh's vertices
	for (UINT i = 0; i < mesh->mNumVertices; i++)
	{
		float x = mesh->mVertices[i].x;
		float y = mesh->mVertices[i].y;
		float z = mesh->mVertices[i].z;
		float nx = 0.0, ny = 0.0, nz = 0.0;
		if (mesh->HasNormals())
		{
			nx = mesh->mNormals[i].x;
			ny = mesh->mNormals[i].y;
			nz = mesh->mNormals[i].z;
		}
		float tx = 0.0, ty = 0.0;
		if (mesh->mTextureCoords[0])
		{
			tx = (float)mesh->mTextureCoords[0][i].x;
			ty = (float)mesh->mTextureCoords[0][i].y;
		}
		Vertex vertex;
		vertex.Pos = XMFLOAT3(x, y, z);
		vertex.Normal = XMFLOAT3(nx, ny, nz);
		vertex.Texcoord = XMFLOAT2(tx, ty);
		vertices.emplace_back(std::move(vertex));
	}
	for (UINT i = 0; i < mesh->mNumFaces; i++)
	{
		aiFace face = mesh->mFaces[i];
		for (UINT j = 0; j < face.mNumIndices; j++)
			indices.push_back(face.mIndices[j]);
	}
	submesh.IndexCount = indices.size() - submesh.StartIndexLocation;
	
	// materials
	std::string material_name = std::to_string(mesh->mMaterialIndex) + "_model_diffuse";
	
	if (mesh->mMaterialIndex >= 0)
	{
		aiMaterial* material = scene->mMaterials[mesh->mMaterialIndex];
		loadMaterialTextures(material, aiTextureType_DIFFUSE, material_name, scene);
		//loadMaterialTextures(material, aiTextureType_SPECULAR, model_name + "_texture_specular", scene);
	}

	auto item = std::make_unique<RenderItem>();
	XMStoreFloat4x4(&item->World, model_matrix);
	XMStoreFloat4x4(&item->TexTransform, XMMatrixScaling(1.0f, 1.0f, 1.0f));
	item->ObjCBIndex = m_objIdx++;
	item->Mat = pRender->m_materials[material_name].get();
	item->Geo = geo.get();
	item->m_primType = D3D_PRIMITIVE_TOPOLOGY_TRIANGLELIST;
	item->indexCount = submesh.IndexCount;
	item->startIndexLocation = submesh.StartIndexLocation;
	item->baseVertexLocation = submesh.BaseVertexLocation;
	pRender->m_allRitems.push_back(std::move(item));
	/**/
	geo->DrawArgs[geo->Name + std::to_string(idx++)] = submesh;
}

void ModelLoader::processNode(aiNode* node, const aiScene* scene)
{
	for (UINT i = 0; i < node->mNumMeshes; i++)
	{
		aiMesh* mesh = scene->mMeshes[node->mMeshes[i]];
		processMesh(mesh, scene);
	}

	for (UINT i = 0; i < node->mNumChildren; i++)
	{
		processNode(node->mChildren[i], scene);
	}
}

void ModelLoader::Close()
{

}

void ModelLoader::loadMaterialTextures(aiMaterial* mat, aiTextureType type, 
	const std::string& typeName, const aiScene* scene)
{
	for (UINT i = 0; i < mat->GetTextureCount(type); i++)
	{
		aiString str;
		mat->GetTexture(type, i, &str);
		bool skip = false;
		for (const auto& x : pRender->m_textures)
		{
			if (x->Name == typeName)
			{
				skip = true;
				break;
			}
		}
		if (!skip)
		{
			// If texture hasn't been loaded already, load it
			if (textype == "embedded compressed texture")
			{
				int textureindex = getTextureIndex(&str);
				getTextureFromModel(scene, textureindex, typeName);
			}
			else
			{
				// load from path;
				std::string filename(str.C_Str());
				filename.replace(filename.end()-3, filename.end(), "dds");
				filename = directory + '/' + filename;
				std::wstring filenamews = std::wstring(filename.begin(), filename.end());

				auto Tex = std::make_unique<Texture>();
				Tex->Name = typeName;
				Tex->Filename = filenamews;

				ThrowIfFailed(DirectX::CreateDDSTextureFromFile12(pRender->m_device.Get(),
					pRender->m_commandList.Get(), Tex->Filename.c_str(),
					Tex->Resource, Tex->UploadHeap));

				auto material = std::make_unique<Material>();
				material->Name = Tex->Name;
				material->MatCBIndex = m_matCBIndex++;
				material->DiffuseSrvHeapIndex = m_texIndex++;
				material->DiffuseAlbedo = XMFLOAT4(1.0f, 1.0f, 1.0f, 1.0f);
				material->FresnelR0 = XMFLOAT3(0.02f, 0.02f, 0.02f);
				material->Roughness = 0.1f;
				
				pRender->m_materials[Tex->Name] = std::move(material);
				pRender->m_textures.push_back(std::move(Tex));
				

				//std::cout << "load texture" << std::endl;
			}
			// Store it as texture loaded for entire model, to ensure we won't unnecesery load duplicate textures.
		}

	}
}

std::string ModelLoader::determineTextureType(const aiScene* scene, aiMaterial* mat)
{
	aiString textypeStr;
	mat->GetTexture(aiTextureType_DIFFUSE, 0, &textypeStr);
	std::string textypeteststr = textypeStr.C_Str();
	if (textypeteststr == "*0" || textypeteststr == "*1" || textypeteststr == "*2" || textypeteststr == "*3" || textypeteststr == "*4" || textypeteststr == "*5")
	{
		if (scene->mTextures[0]->mHeight == 0)
		{
			return "embedded compressed texture";
		}
		else
		{
			return "embedded non-compressed texture";
		}
	}
	if (textypeteststr.find('.') != std::string::npos)
	{
		return "textures are on disk";
	}
}

int ModelLoader::getTextureIndex(aiString* str)
{
	std::string tistr;
	
	tistr = str->C_Str();
	tistr = tistr.substr(1);
	return stoi(tistr);
}

void ModelLoader::getTextureFromModel(const aiScene* scene, int textureindex, const std::string& typeName)
{
	int* size = reinterpret_cast<int*>(&scene->mTextures[textureindex]->mWidth);

	auto Tex = std::make_unique<Texture>();
	Tex->Name = typeName;
	ThrowIfFailed(DirectX::CreateDDSTextureFromMemory12(pRender->m_device.Get(),
		pRender->m_commandList.Get(),
		reinterpret_cast<unsigned char*>(scene->mTextures[textureindex]->pcData),
		*size, Tex->Resource, Tex->UploadHeap));

	auto material = std::make_unique<Material>();
	material->Name = Tex->Name;
	material->MatCBIndex = m_matCBIndex++;
	material->DiffuseSrvHeapIndex = m_texIndex++;
	material->DiffuseAlbedo = XMFLOAT4(1.0f, 1.0f, 1.0f, 1.0f);
	material->FresnelR0 = XMFLOAT3(0.01f, 0.01f, 0.01f);
	material->Roughness = 0.9f;
	
	pRender->m_materials[Tex->Name] = std::move(material);
	pRender->m_textures.push_back(std::move(Tex));
}
